from torch_geometric.utils import scatter
from torch_geometric.datasets import TUDataset
from torch_geometric.loader import DataLoader

dataset = TUDataset(root='/tmp/ENZYMES', name='ENZYMES', use_node_attr=True)
loader = DataLoader(dataset, batch_size=16, shuffle=True)
# 一个batch中抽取15张图的数据，每一张图有多个节点

print(len(dataset))

for data in loader:
    print(len(data.batch))
    print(data.x.shape)
    print(data.y)
    break

